#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文字转语音模块
Text-to-Speech Module

基于讯飞语音API实现文字转语音功能
Text-to-speech functionality based on iFlytek Voice API
"""

import os
import json
import base64
import hmac
import hashlib
import time
import ssl
import wave
import subprocess
import threading
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time
from typing import Optional, Callable, Dict, Any, Tuple
import logging

try:
    import websocket
except ImportError:
    raise ImportError("websocket-client is required. Install with: pip install websocket-client")


class TextToSpeechClient:
    """
    文字转语音客户端
    Text-to-Speech Client
    
    使用讯飞语音API将文字转换为语音
    Convert text to speech using iFlytek Voice API
    """
    
    def __init__(self, app_id: str, api_key: str, api_secret: str):
        """
        初始化TTS客户端
        Initialize TTS client
        
        Args:
            app_id (str): 讯飞应用ID / iFlytek App ID
            api_key (str): API密钥 / API Key
            api_secret (str): API密钥 / API Secret
        """
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        
        # 音频参数
        self.audio_format = "raw"  # 音频格式
        self.audio_encoding = "audio/L16;rate=16000"  # 音频编码
        self.voice_name = "x4_yezi"  # 发音人
        self.text_encoding = "utf8"  # 文本编码
        
        # PCM音频参数
        self.pcm_channels = 1
        self.pcm_sample_width = 2  # 16-bit => 2 bytes
        self.pcm_sample_rate = 16000
        
        # WebSocket连接
        self.ws = None
        self.is_connected = False
        
        # 音频数据缓冲
        self.audio_buffer = bytearray()
        self.synthesis_complete = False
        self.synthesis_error = None
        
        # 回调函数
        self.on_audio_chunk: Optional[Callable[[bytes], None]] = None
        self.on_synthesis_complete: Optional[Callable[[str], None]] = None
        self.on_error: Optional[Callable[[str, int], None]] = None
        
        # 日志
        self.logger = logging.getLogger(__name__)
        
    def set_voice_parameters(self, voice_name: str = "x4_yezi", 
                           audio_format: str = "raw",
                           sample_rate: int = 16000) -> None:
        """
        设置语音参数
        Set voice parameters
        
        Args:
            voice_name (str): 发音人名称 / Voice name
            audio_format (str): 音频格式 / Audio format
            sample_rate (int): 采样率 / Sample rate
        """
        self.voice_name = voice_name
        self.audio_format = audio_format
        self.audio_encoding = f"audio/L16;rate={sample_rate}"
        self.pcm_sample_rate = sample_rate
        
    def set_callbacks(self, 
                     on_audio_chunk: Optional[Callable[[bytes], None]] = None,
                     on_synthesis_complete: Optional[Callable[[str], None]] = None,
                     on_error: Optional[Callable[[str, int], None]] = None) -> None:
        """
        设置回调函数
        Set callback functions
        
        Args:
            on_audio_chunk: 音频块回调 / Audio chunk callback
            on_synthesis_complete: 合成完成回调 / Synthesis complete callback
            on_error: 错误回调 / Error callback
        """
        self.on_audio_chunk = on_audio_chunk
        self.on_synthesis_complete = on_synthesis_complete
        self.on_error = on_error
        
    def _create_url(self) -> str:
        """
        创建WebSocket连接URL
        Create WebSocket connection URL
        
        Returns:
            str: WebSocket URL
        """
        url = 'wss://tts-api.xfyun.cn/v2/tts'
        
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))
        
        # 拼接字符串
        signature_origin = "host: " + "tts-api.xfyun.cn" + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + "/v2/tts " + "HTTP/1.1"
        
        # 进行hmac-sha256加密
        signature_sha = hmac.new(self.api_secret.encode('utf-8'), 
                               signature_origin.encode('utf-8'),
                               digestmod=hashlib.sha256).digest()
        signature_sha = base64.b64encode(signature_sha).decode(encoding='utf-8')
        
        authorization_origin = (f'api_key="{self.api_key}", algorithm="hmac-sha256", '
                              f'headers="host date request-line", signature="{signature_sha}"')
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        
        # 将请求的鉴权参数组合为字典
        params = {
            "authorization": authorization,
            "date": date,
            "host": "tts-api.xfyun.cn"
        }
        
        # 拼接鉴权参数，生成url
        url = url + '?' + urlencode(params)
        return url
        
    def _on_message(self, ws, message):
        """
        WebSocket消息处理
        WebSocket message handler
        """
        try:
            # 检查消息是否为空
            if not message:
                self.logger.warning("Received empty message from TTS WebSocket")
                return
                
            message_data = json.loads(message)
            
            # 安全获取必要字段
            code = message_data.get("code")
            sid = message_data.get("sid")
            
            if code is None:
                self.logger.warning("Received message without code field")
                return
                
            if code != 0:
                error_msg = message_data.get("message", "Unknown error")
                self.synthesis_error = f"TTS API error: {error_msg} (code: {code})"
                self.logger.error(self.synthesis_error)
                if self.on_error:
                    self.on_error(error_msg, code)
                return
                
            # 安全获取数据字段
            data_section = message_data.get("data")
            if not data_section:
                self.logger.warning("Received message without data section")
                return
                
            # 获取音频数据
            audio_data = data_section.get("audio")
            status = data_section.get("status")
            
            if audio_data is None:
                self.logger.warning("Received message without audio data")
                return
                
            if status is None:
                self.logger.warning("Received message without status field")
                return
                
            try:
                audio_bytes = base64.b64decode(audio_data)
            except Exception as decode_error:
                self.logger.error(f"Failed to decode audio data: {decode_error}")
                return
            
            # 添加到缓冲区
            self.audio_buffer.extend(audio_bytes)
            
            # 调用音频块回调
            if self.on_audio_chunk:
                try:
                    self.on_audio_chunk(audio_bytes)
                except Exception as callback_error:
                    self.logger.error(f"Error in audio chunk callback: {callback_error}")
                
            # 检查是否为最后一帧
            if status == 2:  # 最后一帧
                self.synthesis_complete = True
                self.logger.info(f"TTS synthesis completed, total audio bytes: {len(self.audio_buffer)}")
                try:
                    ws.close()
                except Exception as close_error:
                    self.logger.warning(f"Error closing WebSocket: {close_error}")
                
        except json.JSONDecodeError as json_error:
            self.synthesis_error = f"Failed to parse TTS message JSON: {str(json_error)}"
            self.logger.error(self.synthesis_error)
            if self.on_error:
                self.on_error(str(json_error), -1)
        except Exception as e:
            self.synthesis_error = f"Failed to process TTS message: {str(e)}"
            self.logger.error(self.synthesis_error)
            if self.on_error:
                self.on_error(str(e), -1)
                
    def _on_error(self, ws, error):
        """
        WebSocket错误处理
        WebSocket error handler
        """
        self.synthesis_error = f"WebSocket error: {str(error)}"
        self.logger.error(self.synthesis_error)
        if self.on_error:
            self.on_error(str(error), -1)
            
    def _on_close(self, ws, close_status_code, close_msg):
        """
        WebSocket关闭处理
        WebSocket close handler
        """
        self.is_connected = False
        self.logger.info("TTS WebSocket connection closed")
        
    def _on_open(self, ws):
        """
        WebSocket连接打开处理
        WebSocket open handler
        """
        self.is_connected = True
        self.logger.info("TTS WebSocket connection opened")
        
    def synthesize_text(self, text: str, output_file: Optional[str] = None, max_retries: int = 3) -> Tuple[bool, str, Optional[bytes]]:
        """
        合成文字为语音
        Synthesize text to speech
        
        Args:
            text (str): 要合成的文字 / Text to synthesize
            output_file (Optional[str]): 输出文件路径 / Output file path
            max_retries (int): 最大重试次数 / Maximum retry attempts
            
        Returns:
            Tuple[bool, str, Optional[bytes]]: (成功标志, 消息, 音频数据)
        """
        # 输入验证
        if not text or not text.strip():
            return False, "Text cannot be empty", None
            
        if len(text) > 8000:  # 讯飞TTS文本长度限制
            return False, "Text too long (max 8000 characters)", None
            
        # 重试逻辑
        for attempt in range(max_retries):
            try:
                self.logger.info(f"TTS synthesis attempt {attempt + 1}/{max_retries} for text: {text[:50]}...")
                
                # 重置状态
                self.audio_buffer = bytearray()
                self.synthesis_complete = False
                self.synthesis_error = None
                
                # 创建WebSocket连接
                ws_url = self._create_url()
                self.ws = websocket.WebSocketApp(
                    ws_url,
                    on_message=self._on_message,
                    on_error=self._on_error,
                    on_close=self._on_close,
                    on_open=self._on_open
                )
            
                # 准备发送数据
                def send_data():
                    try:
                        # 构建请求数据
                        data = {
                            "common": {"app_id": self.app_id},
                            "business": {
                                "aue": self.audio_format,
                                "auf": self.audio_encoding,
                                "vcn": self.voice_name,
                                "tte": self.text_encoding
                            },
                            "data": {
                                "status": 2,
                                "text": str(base64.b64encode(text.encode('utf-8')), "UTF8")
                            }
                        }
                        
                        # 发送数据
                        self.ws.send(json.dumps(data))
                        self.logger.debug("TTS data sent successfully")
                    except Exception as send_error:
                        self.synthesis_error = f"Failed to send TTS data: {str(send_error)}"
                        self.logger.error(self.synthesis_error)
                    
                # 设置发送数据回调
                self.ws.on_open = lambda ws: send_data()
                
                # 在单独线程中运行WebSocket，避免阻塞主线程
                import threading
                ws_thread = threading.Thread(
                    target=lambda: self.ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
                )
                ws_thread.daemon = True
                ws_thread.start()
                
                # 等待合成完成，但设置超时避免无限等待
                timeout = 30  # 30秒超时
                start_time = time.time()
                
                while not self.synthesis_complete and not self.synthesis_error:
                    if time.time() - start_time > timeout:
                        self.synthesis_error = "TTS synthesis timeout"
                        break
                    time.sleep(0.1)  # 短暂休眠，避免CPU占用过高
                
                # 强制关闭WebSocket连接
                try:
                    if self.ws:
                        self.ws.close()
                except Exception as close_error:
                    self.logger.warning(f"Error closing WebSocket: {close_error}")
                
                # 等待线程结束
                ws_thread.join(timeout=2)
                
                # 检查合成结果
                if self.synthesis_error:
                    if attempt < max_retries - 1:  # 不是最后一次尝试
                        self.logger.warning(f"TTS attempt {attempt + 1} failed: {self.synthesis_error}, retrying...")
                        time.sleep(1)  # 重试前等待1秒
                        continue
                    else:
                        return False, f"TTS failed after {max_retries} attempts: {self.synthesis_error}", None
                    
                if not self.synthesis_complete or len(self.audio_buffer) == 0:
                    if attempt < max_retries - 1:  # 不是最后一次尝试
                        self.logger.warning(f"TTS attempt {attempt + 1} incomplete, retrying...")
                        time.sleep(1)  # 重试前等待1秒
                        continue
                    else:
                        return False, f"TTS synthesis failed or no audio data received after {max_retries} attempts", None
                    
                # 成功合成
                self.logger.info(f"TTS synthesis successful on attempt {attempt + 1}")
                
                # 保存到文件（如果指定）
                if output_file:
                    success, message = self._save_audio_to_file(bytes(self.audio_buffer), output_file)
                    if not success:
                        return False, message, bytes(self.audio_buffer)
                        
                return True, f"TTS synthesis completed successfully, audio bytes: {len(self.audio_buffer)}", bytes(self.audio_buffer)
                
            except Exception as e:
                error_msg = f"TTS attempt {attempt + 1} failed: {str(e)}"
                self.logger.error(error_msg)
                if attempt < max_retries - 1:  # 不是最后一次尝试
                    self.logger.warning(f"Retrying TTS synthesis...")
                    time.sleep(1)  # 重试前等待1秒
                    continue
                else:
                    return False, f"TTS failed after {max_retries} attempts: {str(e)}", None
        
        # 如果所有重试都失败了（理论上不会到达这里）
        return False, f"TTS synthesis failed after {max_retries} attempts", None
            
    def _save_audio_to_file(self, audio_data: bytes, output_file: str) -> Tuple[bool, str]:
        """
        保存音频数据到文件
        Save audio data to file
        
        Args:
            audio_data (bytes): 音频数据 / Audio data
            output_file (str): 输出文件路径 / Output file path
            
        Returns:
            Tuple[bool, str]: (成功标志, 消息)
        """
        try:
            # 确定文件格式
            file_ext = os.path.splitext(output_file)[1].lower()
            
            if file_ext == '.pcm':
                # 直接保存PCM数据
                with open(output_file, 'wb') as f:
                    f.write(audio_data)
                return True, f"PCM audio saved to {output_file}"
                
            elif file_ext == '.wav':
                # 转换为WAV格式
                return self._pcm_to_wav(audio_data, output_file)
                
            else:
                return False, f"Unsupported audio format: {file_ext}"
                
        except Exception as e:
            error_msg = f"Failed to save audio file: {str(e)}"
            self.logger.error(error_msg)
            return False, error_msg
            
    def _pcm_to_wav(self, pcm_data: bytes, wav_path: str) -> Tuple[bool, str]:
        """
        将PCM数据转换为WAV格式
        Convert PCM data to WAV format
        
        Args:
            pcm_data (bytes): PCM音频数据 / PCM audio data
            wav_path (str): WAV文件路径 / WAV file path
            
        Returns:
            Tuple[bool, str]: (成功标志, 消息)
        """
        try:
            # 首先尝试使用wave库
            with wave.open(wav_path, 'wb') as wav_file:
                wav_file.setnchannels(self.pcm_channels)
                wav_file.setsampwidth(self.pcm_sample_width)
                wav_file.setframerate(self.pcm_sample_rate)
                wav_file.writeframes(pcm_data)
                
            return True, f"WAV audio saved to {wav_path}"
            
        except Exception as e:
            # 如果wave库失败，尝试使用ffmpeg
            self.logger.warning(f"Wave library failed: {e}, trying ffmpeg...")
            return self._pcm_to_wav_with_ffmpeg(pcm_data, wav_path)
            
    def _pcm_to_wav_with_ffmpeg(self, pcm_data: bytes, wav_path: str) -> Tuple[bool, str]:
        """
        使用ffmpeg将PCM数据转换为WAV格式
        Convert PCM data to WAV format using ffmpeg
        
        Args:
            pcm_data (bytes): PCM音频数据 / PCM audio data
            wav_path (str): WAV文件路径 / WAV file path
            
        Returns:
            Tuple[bool, str]: (成功标志, 消息)
        """
        try:
            # 创建临时PCM文件
            temp_pcm = wav_path.replace('.wav', '_temp.pcm')
            
            # 写入PCM数据
            with open(temp_pcm, 'wb') as f:
                f.write(pcm_data)
                
            # 使用ffmpeg转换
            cmd = [
                'ffmpeg', '-y',
                '-f', 's16le',
                '-ar', str(self.pcm_sample_rate),
                '-ac', str(self.pcm_channels),
                '-i', temp_pcm,
                wav_path
            ]
            
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            # 清理临时文件
            if os.path.exists(temp_pcm):
                os.remove(temp_pcm)
                
            if result.returncode == 0:
                return True, f"WAV audio converted and saved to {wav_path}"
            else:
                return False, f"ffmpeg conversion failed: {result.stderr}"
                
        except Exception as e:
            error_msg = f"Failed to convert PCM to WAV with ffmpeg: {str(e)}"
            self.logger.error(error_msg)
            return False, error_msg